#include <iostream>
#include <fstream>
#include <vector>
#include <chrono>
#include <string>
#include <iomanip>
#include <cstring>
#include <fcntl.h>
#include <unistd.h>
#include <sys/stat.h>
#include <sys/statvfs.h>
#include <signal.h>
#include <atomic>

// 全局变量控制程序退出
std::atomic<bool> g_running(true);

// 信号处理函数
void signal_handler(int signal) {
    if (signal == SIGINT || signal == SIGTERM) {
        std::cout << "\n收到退出信号，正在安全退出..." << std::endl;
        g_running = false;
    }
}

// 基于文件大小的写入测试配置
struct SizeBasedTestConfig {
    std::string path;           // 测试文件路径
    size_t block_size_kb;       // 写入块大小(KB) - 代码中指定
    size_t target_file_size_mb; // 目标文件大小(MB)
    bool use_direct_io;         // 是否使用O_DIRECT
    bool use_o_sync;            // 是否使用O_SYNC
    size_t report_interval_mb;  // 报告间隔(MB)
};

// 单位转换辅助函数
size_t parse_size_with_unit(const std::string& size_str) {
    if (size_str.empty()) return 0;

    std::string num_part;
    std::string unit_part;

    // 分离数字和单位
    size_t i = 0;
    while (i < size_str.length() && (isdigit(size_str[i]) || size_str[i] == '.')) {
        num_part += size_str[i];
        i++;
    }

    while (i < size_str.length()) {
        unit_part += tolower(size_str[i]);
        i++;
    }

    double num = std::stod(num_part);

    // 转换为MB
    if (unit_part.empty() || unit_part == "mb" || unit_part == "m") {
        return static_cast<size_t>(num);
    } else if (unit_part == "gb" || unit_part == "g") {
        return static_cast<size_t>(num * 1024);
    } else if (unit_part == "kb" || unit_part == "k") {
        return static_cast<size_t>(num / 1024.0);
    } else if (unit_part == "tb" || unit_part == "t") {
        return static_cast<size_t>(num * 1024 * 1024);
    } else {
        std::cerr << "不支持的单位: " << unit_part << std::endl;
        return 0;
    }
}

// 格式化文件大小显示
std::string format_size(size_t size_mb) {
    if (size_mb >= 1024 * 1024) {
        double tb = size_mb / (1024.0 * 1024.0);
        return std::to_string(tb).substr(0, 4) + " TB";
    } else if (size_mb >= 1024) {
        double gb = size_mb / 1024.0;
        return std::to_string(gb).substr(0, 5) + " GB";
    } else {
        return std::to_string(size_mb) + " MB";
    }
}

// 基于文件大小的写入测试类
class SizeBasedWriteTest {
private:
    SizeBasedTestConfig config_;
    int fd_;
    void* buffer_;
    size_t buffer_size_;

    // 统计信息
    size_t total_bytes_written_;
    size_t bytes_since_last_report_;
    size_t target_bytes_;
    std::chrono::high_resolution_clock::time_point start_time_;
    std::chrono::high_resolution_clock::time_point last_report_time_;

public:
    SizeBasedWriteTest(const SizeBasedTestConfig& config)
            : config_(config), fd_(-1), buffer_(nullptr),
              total_bytes_written_(0), bytes_since_last_report_(0) {

        buffer_size_ = config_.block_size_kb * 1024;
        target_bytes_ = config_.target_file_size_mb * 1024 * 1024;
    }

    ~SizeBasedWriteTest() {
        cleanup();
    }

    bool initialize() {
        // 创建测试文件
        std::string filename = config_.path + "/size_based_write_test.dat";

        int flags = O_WRONLY | O_CREAT | O_TRUNC;
        if (config_.use_direct_io) {
            flags |= O_DIRECT;
        }
        if (config_.use_o_sync) {
            flags |= O_SYNC;
        }

        fd_ = open(filename.c_str(), flags, 0644);
        if (fd_ < 0) {
            std::cerr << "无法打开文件: " << filename << " 错误: " << strerror(errno) << std::endl;
            return false;
        }

        // 分配缓冲区
        if (config_.use_direct_io) {
            // DirectIO需要内存对齐
            if (posix_memalign(&buffer_, 4096, buffer_size_) != 0) {
                std::cerr << "内存对齐分配失败" << std::endl;
                return false;
            }
        } else {
            buffer_ = malloc(buffer_size_);
            if (!buffer_) {
                std::cerr << "内存分配失败" << std::endl;
                return false;
            }
        }

        // 填充测试数据 (使用不同的模式以避免压缩优化)
        char* char_buffer = (char*)buffer_;
        for (size_t i = 0; i < buffer_size_; i++) {
            char_buffer[i] = (char)(i % 256);
        }

        return true;
    }

    void cleanup() {
        if (fd_ >= 0) {
            close(fd_);
            fd_ = -1;
        }
        if (buffer_) {
            free(buffer_);
            buffer_ = nullptr;
        }

        // 删除测试文件
        std::string filename = config_.path + "/size_based_write_test.dat";
        unlink(filename.c_str());
    }

    void print_header() {
        std::cout << "========== 基于文件大小的写入测试 ==========" << std::endl;
        std::cout << "测试配置:" << std::endl;
        std::cout << "  - 测试路径: " << config_.path << std::endl;
        std::cout << "  - 块大小: " << config_.block_size_kb << " KB" << std::endl;
        std::cout << "  - 目标文件大小: " << format_size(config_.target_file_size_mb) << std::endl;
        std::cout << "  - DirectIO: " << (config_.use_direct_io ? "启用" : "禁用") << std::endl;
        std::cout << "  - SyncIO: " << (config_.use_o_sync ? "启用" : "禁用") << std::endl;
        std::cout << "  - 报告间隔: " << config_.report_interval_mb << " MB" << std::endl;
        std::cout << "=============================================" << std::endl;
        std::cout << std::left << std::setw(12) << "已写入(MB)"
                  << std::setw(10) << "进度(%)"
                  << std::setw(15) << "瞬时速度(MB/s)"
                  << std::setw(15) << "平均速度(MB/s)"
                  << std::setw(12) << "剩余时间" << std::endl;
        std::cout << "-------------------------------------------------------------------" << std::endl;
    }

    void print_progress() {
        auto current_time = std::chrono::high_resolution_clock::now();
        size_t written_mb = total_bytes_written_ / (1024 * 1024);

        if (bytes_since_last_report_ >= config_.report_interval_mb * 1024 * 1024) {
            auto duration_since_last = std::chrono::duration_cast<std::chrono::milliseconds>(
                    current_time - last_report_time_).count();

            auto total_duration = std::chrono::duration_cast<std::chrono::milliseconds>(
                    current_time - start_time_).count();

            if (duration_since_last > 0 && total_duration > 0) {
                double instant_speed = (bytes_since_last_report_ / (1024.0 * 1024.0)) /
                                       (duration_since_last / 1000.0);
                double average_speed = (total_bytes_written_ / (1024.0 * 1024.0)) /
                                       (total_duration / 1000.0);

                double progress = (double)total_bytes_written_ / target_bytes_ * 100.0;

                // 计算剩余时间
                std::string eta_str = "未知";
                if (average_speed > 0) {
                    double remaining_mb = (target_bytes_ - total_bytes_written_) / (1024.0 * 1024.0);
                    double eta_seconds = remaining_mb / average_speed;

                    if (eta_seconds < 60) {
                        eta_str = std::to_string((int)eta_seconds) + "s";
                    } else if (eta_seconds < 3600) {
                        eta_str = std::to_string((int)(eta_seconds / 60)) + "m" +
                                  std::to_string((int)eta_seconds % 60) + "s";
                    } else {
                        int hours = (int)(eta_seconds / 3600);
                        int minutes = (int)((eta_seconds - hours * 3600) / 60);
                        eta_str = std::to_string(hours) + "h" + std::to_string(minutes) + "m";
                    }
                }

                std::cout << std::left << std::setw(12) << written_mb
                          << std::setw(10) << std::fixed << std::setprecision(1) << progress
                          << std::setw(15) << std::fixed << std::setprecision(2) << instant_speed
                          << std::setw(15) << std::fixed << std::setprecision(2) << average_speed
                          << std::setw(12) << eta_str
                          << std::endl;
            }

            // 重置统计
            bytes_since_last_report_ = 0;
            last_report_time_ = current_time;
        }
    }

    void print_summary() {
        auto end_time = std::chrono::high_resolution_clock::now();
        auto total_duration = std::chrono::duration_cast<std::chrono::milliseconds>(
                end_time - start_time_).count();

        double total_seconds = total_duration / 1000.0;
        double total_mb = total_bytes_written_ / (1024.0 * 1024.0);
        double average_speed = total_mb / total_seconds;
        double progress = (double)total_bytes_written_ / target_bytes_ * 100.0;

        std::cout << "\n========== 测试完成 ==========" << std::endl;
        std::cout << "目标文件大小: " << format_size(config_.target_file_size_mb) << std::endl;
        std::cout << "实际写入大小: " << std::fixed << std::setprecision(2) << total_mb << " MB" << std::endl;
        std::cout << "完成进度: " << std::fixed << std::setprecision(1) << progress << "%" << std::endl;
        std::cout << "总耗时: " << std::fixed << std::setprecision(2) << total_seconds << " 秒" << std::endl;
        std::cout << "平均写入速度: " << std::fixed << std::setprecision(2) << average_speed << " MB/s" << std::endl;
        std::cout << "===============================" << std::endl;
    }

    bool run() {
        if (!initialize()) {
            return false;
        }

        print_header();

        start_time_ = std::chrono::high_resolution_clock::now();
        last_report_time_ = start_time_;

        while (g_running && total_bytes_written_ < target_bytes_) {
            // 计算本次写入大小
            size_t remaining_bytes = target_bytes_ - total_bytes_written_;
            size_t write_size = std::min(buffer_size_, remaining_bytes);

            // 如果使用DirectIO，确保写入大小对齐
            if (config_.use_direct_io && write_size % 512 != 0) {
                write_size = (write_size / 512) * 512;
                if (write_size == 0) break; // 剩余数据太少，无法对齐
            }

            // 执行写入操作
            ssize_t written = write(fd_, buffer_, write_size);
            if (written < 0) {
                std::cerr << "\n写入错误: " << strerror(errno) << std::endl;
                break;
            }

            total_bytes_written_ += written;
            bytes_since_last_report_ += written;

            // 定期打印进度
            print_progress();

            // 如果不使用O_SYNC，可以定期调用fsync确保数据写入
            if (!config_.use_o_sync && (total_bytes_written_ % (100 * 1024 * 1024)) == 0) {
                fsync(fd_);
            }
        }

        // 确保所有数据写入磁盘
        if (fd_ >= 0) {
            fsync(fd_);
        }

        print_summary();
        return true;
    }
};

void print_usage(const char* program_name) {
    std::cout << "用法: " << program_name << " <测试目录路径> <文件大小> [选项]" << std::endl;
    std::cout << std::endl;
    std::cout << "文件大小格式:" << std::endl;
    std::cout << "  数字 + 单位 (K/KB, M/MB, G/GB, T/TB)" << std::endl;
    std::cout << "  例如: 100MB, 1.5GB, 2T, 512M" << std::endl;
    std::cout << std::endl;
    std::cout << "选项:" << std::endl;
    std::cout << "  --direct=<0|1>         是否使用DirectIO (默认: 1)" << std::endl;
    std::cout << "  --sync=<0|1>           是否使用SyncIO (默认: 0)" << std::endl;
    std::cout << "  --report=<MB>          报告间隔，单位MB (默认: 100)" << std::endl;
    std::cout << std::endl;
    std::cout << "示例:" << std::endl;
    std::cout << "  " << program_name << " /mnt/nvme 1GB                    # 写入1GB文件" << std::endl;
    std::cout << "  " << program_name << " ./ 500MB --direct=0              # 写入500MB，禁用DirectIO" << std::endl;
    std::cout << "  " << program_name << " /tmp 2.5GB --report=50           # 写入2.5GB，每50MB报告一次" << std::endl;
    std::cout << "  " << program_name << " /mnt/ssd 10GB --sync=1           # 写入10GB，启用SyncIO" << std::endl;
}

int main(int argc, char* argv[]) {
    std::cout << "基于文件大小的存储写入性能测试工具 v2.0" << std::endl;

    if (argc < 3) {
        print_usage(argv[0]);
        return 1;
    }

    // 设置信号处理
    signal(SIGINT, signal_handler);
    signal(SIGTERM, signal_handler);

    SizeBasedTestConfig config;
    config.path = argv[1];
    config.target_file_size_mb = parse_size_with_unit(argv[2]);
    config.block_size_kb = 1024;        // 固定块大小1MB
    config.use_direct_io = true;        // 默认启用DirectIO
    config.use_o_sync = false;          // 默认禁用SyncIO
    config.report_interval_mb = 100;    // 默认每100MB报告一次

    if (config.target_file_size_mb == 0) {
        std::cerr << "错误: 无效的文件大小格式: " << argv[2] << std::endl;
        print_usage(argv[0]);
        return 1;
    }

    // 解析可选参数
    for (int i = 3; i < argc; i++) {
        std::string arg = argv[i];

        if (arg.find("--direct=") == 0) {
            config.use_direct_io = (std::stoi(arg.substr(9)) != 0);
        } else if (arg.find("--sync=") == 0) {
            config.use_o_sync = (std::stoi(arg.substr(7)) != 0);
        } else if (arg.find("--report=") == 0) {
            config.report_interval_mb = std::stoi(arg.substr(9));
        }
    }

    // 验证参数
    if (config.target_file_size_mb <= 0) {
        std::cerr << "错误: 文件大小必须大于0" << std::endl;
        return 1;
    }

    if (config.report_interval_mb <= 0) {
        std::cerr << "错误: 报告间隔必须大于0" << std::endl;
        return 1;
    }

    // 检查目录
    struct stat st;
    if (stat(config.path.c_str(), &st) != 0 || !S_ISDIR(st.st_mode)) {
        std::cerr << "错误: 指定的路径不是有效目录: " << config.path << std::endl;
        return 1;
    }

    if (access(config.path.c_str(), W_OK) != 0) {
        std::cerr << "错误: 没有写入权限: " << config.path << std::endl;
        return 1;
    }

    // 检查可用空间
    struct statvfs vfs;
    if (statvfs(config.path.c_str(), &vfs) == 0) {
        size_t available_mb = (vfs.f_bavail * vfs.f_frsize) / (1024 * 1024);
        if (available_mb < config.target_file_size_mb) {
            std::cerr << "警告: 可用空间(" << available_mb << " MB) 小于目标文件大小("
                      << config.target_file_size_mb << " MB)" << std::endl;
            std::cout << "是否继续? (y/N): ";
            char choice;
            std::cin >> choice;
            if (choice != 'y' && choice != 'Y') {
                return 1;
            }
        }
    }

    // 运行测试
    SizeBasedWriteTest test(config);
    bool success = test.run();

    return success ? 0 : 1;
}